package com.itheima.ai.config;

import jakarta.annotation.PostConstruct;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.CommandLineRunner;
import org.springframework.stereotype.Component;
import java.util.List;

@Component
public class RagInitializer  {

    @Autowired
    private PagePdfDocumentReader pdfReader;
    @Autowired
    private VectorStore vectorStore;
    @Autowired
    private OllamaEmbeddingModel ollamaEmbeddingModel;

    @PostConstruct
    public void initializePdfToVectorStore() {
        try {
            List<Document> documents = pdfReader.read();
            System.out.println("PDF页数：" + documents.size());

            // 打印第一条文档的向量维度（验证模型维度）
            if (!documents.isEmpty()) {
                Document firstDoc = documents.get(0);
                float[] embedding = ollamaEmbeddingModel.embed(firstDoc.getText());
                System.out.println("模型生成的向量维度：" + embedding.length); // 应输出4096
            }

            vectorStore.add(documents);
            System.out.println("数据写入成功");
        } catch (Exception e) {
            e.printStackTrace();
        }
    }


//    @Override
//    public void run(String... args) throws Exception {
//        // 读取PDF并写入向量库
//        List<Document> documents = pdfReader.read();
//        System.out.println("成功读取PDF，共" + documents.size() + "页");
//        // 打印第一条文档的向量维度（验证模型维度）
//        if (!documents.isEmpty()) {
//            Document firstDoc = documents.get(0);
//            float[] embedding = ollamaEmbeddingModel.embed(firstDoc.getContent()).getResult();
//            System.out.println("模型生成的向量维度：" + embedding.length); // 应输出4096
//        }
//
//        vectorStore.add(documents);
//        System.out.println("PDF内容已成功写入向量库");
//    }
}
